import rospy
from nav_msgs.msg import Odometry
from geometry_msgs.msg import PoseStamped
from nav_msgs.msg import Path
from mavros_msgs.msg import PositionTarget

#!/usr/bin/env python

class TrajectoryVisualizer:
    def __init__(self):
        rospy.init_node('trajectory_visualizer', anonymous=True)
        
        # Create a publisher for the path
        self.path_pub = rospy.Publisher('/iris_0/trajectory', Path,queue_size=1000)
        
        # Subscribe to the odometry topic
        self.odom_sub = rospy.Subscriber('/iris_0/transfer/Odometry', Odometry, self.odom_callback,queue_size=1000)
        
        # Subscribe to the local setpoint topic
        self.local_setpoint_sub = rospy.Subscriber('/iris_0/mavros/setpoint_raw/local', PositionTarget, self.local_setpoint_callback,queue_size=1000)
        
        # Initialize the paths message
        self.path = Path()
        self.path.header.frame_id = "world"  # Set the frame ID
        
        self.setpoint_path = Path()
        self.setpoint_path.header.frame_id = "world"
        self.setpoint_path_pub = rospy.Publisher('/iris_0/setpoint_trajectory', Path, queue_size=1000)

    def odom_callback(self, msg):
        # Create a PoseStamped message
        pose = PoseStamped()
        pose.header = msg.header
        pose.pose = msg.pose.pose
        
        # Append the new pose to the path
        self.path.poses.append(pose)
        
        # Publish the updated path
        self.path_pub.publish(self.path)

    def local_setpoint_callback(self, msg):
        # Create a PoseStamped message for the setpoint
        pose = PoseStamped()
        pose.header = msg.header
        pose.pose.position.x = msg.position.x  # Only use position (x,y,z)
        pose.pose.position.y = msg.position.y
        pose.pose.position.z = msg.position.z
        # Append the new pose to the setpoint path
        self.setpoint_path.poses.append(pose)
        
        # Publish the updated setpoint path
        self.setpoint_path_pub.publish(self.setpoint_path)

if __name__ == '__main__':
    try:
        visualizer = TrajectoryVisualizer()
        rospy.spin()
    except rospy.ROSInterruptException:
        pass